import contextlib
import functools
import itertools
import threading


class _State(threading.local):
    def __init__(self):
        self.contexts = ()


_state = _State()


@contextlib.contextmanager
def StackContext(context_factory):
    old_contexts = _state.contexts
    try:
        _state.contexts = old_contexts + (context_factory,)
        with context_factory():
            yield
    finally:
        _state.contexts = old_contexts


@contextlib.contextmanager
def NullContext():
    old_contexts = _state.contexts
    try:
        _state.contexts = ()
        yield
    finally:
        _state.contexts = old_contexts


def wrap(fn):
    def wrapped(callback, contexts, *args, **kwargs):
        pairs = itertools.izip(itertools.chain(_state.contexts,
                                               itertools.repeat(None)),
                               contexts)
        new_contexts = []
        for old, new in itertools.dropwhile(lambda x: x[0] is x[1], pairs):
            new_contexts.append(StackContext(new))
        if new_contexts:
            with contextlib.nested(*new_contexts):
                callback(*args, **kwargs)
        else:
            callback(*args, **kwargs)
    if getattr(fn, 'stack_context_wrapped', False):
        return fn
    contexts = _state.contexts
    result = functools.partial(wrapped, fn, contexts)
    result.stack_context_wrapped = True
    return result
